import os
import torch
import torch.utils.data
import torchvision
import numpy as np

from data.apple_dataset import AppleDataset
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

import utility.utils as utils
import utility.transforms as T


######################################################
# 使用 MinneApple 数据集，通过 Faster-RCNN 或 Mask-RCNN 预测器进行预测
######################################################
def get_transform(train):
    # 定义数据转换流程
    transforms = []
    transforms.append(T.ToTensor())  # 将图像转换为张量
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))  # 随机水平翻转，概率为 0.5
    return T.Compose(transforms)  # 组合所有转换


def get_maskrcnn_model_instance(num_classes):
    # 加载预训练的 Mask R-CNN 模型
    # 加载在 COCO 数据集上预训练的实例分割模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False)

    # 获取分类器的输入特征数量
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 替换预训练的头部为新的头部
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 获取掩码分类器的输入特征数量
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # 替换掩码预测器为新的预测器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    return model


def get_frcnn_model_instance(num_classes):
    # 加载预训练的 Faster R-CNN 模型
    # 加载在 COCO 数据集上预训练的实例分割模型
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)

    # 获取分类器的输入特征数量
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 替换预训练的头部为新的头部
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model


def main(args):
    # 主函数，用于加载模型、数据集，并进行预测
    num_classes = 2  # 类别数量（背景 + 苹果）
    device = args.device  # 使用的设备（CPU 或 GPU）

    # 加载模型
    print("加载模型")
    # 根据参数创建正确的模型类型
    if args.mrcnn:
        model = get_maskrcnn_model_instance(num_classes)
    else:
        model = get_frcnn_model_instance(num_classes)

    # 加载模型参数并保持在 CPU 上
    checkpoint = torch.load(args.weight_file, map_location=device)
    model.load_state_dict(checkpoint['model'], strict=False)
    model.eval()

    print("创建数据加载器")
    dataset_test = AppleDataset(args.data_path, get_transform(train=False))
    data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1,
                                                   shuffle=False, num_workers=1,
                                                   collate_fn=utils.collate_fn)

    # 创建输出目录
    base_path = os.path.dirname(args.output_file)
    if not os.path.exists(base_path):
        os.makedirs(base_path)

    # 对每张图像的边界框进行预测
    f = open(args.output_file, 'a')
    for image, targets in data_loader_test:
        image = list(img.to(device) for img in image)
        outputs = model(image)
        for ii, output in enumerate(outputs):
            img_id = targets[ii]['image_id']
            img_name = data_loader_test.dataset.get_img_name(img_id)
            print("正在预测图像: {}".format(img_name))
            boxes = output['boxes'].detach().numpy()
            scores = output['scores'].detach().numpy()

            im_names = np.repeat(img_name, len(boxes), axis=0)
            stacked = np.hstack((im_names.reshape(len(scores), 1), boxes.astype(int), scores.reshape(len(scores), 1)))

            # 将预测结果写入文件
            np.savetxt(f, stacked, fmt='%s', delimiter=',', newline='\n')


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='PyTorch Detection')
    parser.add_argument('--data_path', required=True, help='要进行预测的数据路径')
    parser.add_argument('--output_file', required=True, help='预测输出结果的保存路径')
    parser.add_argument('--weight_file', required=True, help='权重文件路径')
    parser.add_argument('--device', default='cpu', help='使用的设备。可以是 cpu 或 cuda')
    model = parser.add_mutually_exclusive_group(required=True)
    model.add_argument('--frcnn', action='store_true', help='使用 Faster-RCNN 模型')
    model.add_argument('--mrcnn', action='store_true', help='使用 Mask-RCNN 模型')

    args = parser.parse_args()
    main(args)